import torch
import torch.nn as nn


class Inception_Inverted_ResNet(nn.Module):  ## ！！！ 整体运行后， 输入和输出的图片形状一致
    def __init__(self, in_channels, mid_channels, scale=1.0):
        super(Inception_Inverted_ResNet, self).__init__()
        self.scale = scale
        self.branch_0 = Conv2d(in_channels, mid_channels, 1, stride=1, padding=0, bias=False)  ## 1*1卷积 升维
        self.branch_1 = nn.Sequential(
            Conv2d(in_channels, mid_channels, 1, stride=1, padding=0, bias=False),  ## 1*1卷积
            Conv2d(mid_channels, mid_channels, 3, stride=1, padding=1, bias=False)  ## 代码有误！！！！ 3*3 DW卷积  通道数不变，图片大小正常计算
        )
        self.branch_2 = nn.Sequential(
            Conv2d(in_channels, mid_channels, 1, stride=1, padding=0, bias=False),  ## 1*1卷积
            Conv2d(mid_channels, mid_channels, 5, stride=1, padding=2, bias=False)  ## 5*5 DW卷积
        )
        self.branch_3 = nn.Sequential(
            Conv2d(in_channels, mid_channels, 1, stride=1, padding=0, bias=False),  ## 1*1卷积
            Conv2d(mid_channels, mid_channels, 3, stride=1, padding=1, bias=False),  ## 3*3 DW卷积
            Conv2d(mid_channels, mid_channels, 5, stride=1, padding=2, bias=False)  ## 5*5 DW卷积
        )
        self.conv = nn.Conv2d(4 * mid_channels, in_channels, 1, stride=1, padding=0, bias=True)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x0 = self.branch_0(x)
        x1 = self.branch_1(x)
        x2 = self.branch_2(x)
        x3 = self.branch_3(x)
        x_combine = torch.cat((x0, x1, x2, x3), dim=1)
        # print(x.shape)
        x_out = self.conv(x_combine)
        # print(x_out.shape)
        return self.relu(x + self.scale * x_out)  ## 残差连接


class Conv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding, stride=1, bias=True):
        super(Conv2d, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=bias)
        self.bn = nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.1)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x


if __name__ == '__main__':
    image = torch.Tensor(2, 3072, 576, 576).cuda()
    # out = TriTransUnet(image)
    # Full_Rnet = models.resnet50(pretrained=True)
    net = Inception_Inverted_ResNet(in_channels=3072, mid_channels=4096).cuda()
    # net = Inception_Inverted_ResNet(in_channels=768, mid_channels=4)
    out = net(image)
    print("你好呀")
    print(out.shape)
